#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/3/6 16:13
# @Author  : wanghaoran
# @File    : rerank.py
from advanced_module.base import BaseRerankEngineTool
import json
from configs import device
from reranker import reranker_model


class SingleReranker(BaseRerankEngineTool):
    def __init__(self):
        pass

    def get_rerank_score(self, question, retrieval_chunk):
        rerank_inputs = reranker_model.tokenize(question, retrieval_chunk, return_tensors='pt')
        rerank_inputs = rerank_inputs.to(device)
        rerank_score = reranker_model(rerank_inputs).logits
        return rerank_score.cpu().detach().numpy()[0][0]

    # def get_rerank_score(self, pairs):
    #     encoded_input = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt')
    #
    #     scores_ort = reranker_model(**encoded_input, return_dict=True).logits.view(-1, ).float()
    #     scores = scores_ort.detach().numpy().tolist()
    #     return [round(s, 4) for s in scores]

    def get_results(self, question, candidate_results, topk):
        # pair_data = [[question, c] for c in candidate_results["text"].tolist()]
        # candidate_results["rerank_score"] = self.get_rerank_score(pair_data)
        # print(candidate_results)
        candidate_results["rerank_score"] = candidate_results.text.apply(lambda x: round(self.get_rerank_score(question, x), 4))
        candidate_results.sort_values(by='rerank_score', ascending=False, inplace=True)
        candidate_results = candidate_results.head(topk)
        candidate_results.drop('uid', axis=1, inplace=True)
        output_json = candidate_results.to_dict(orient='records')
        return output_json